#include <cassert>
#include <cstdint>
#include <memory>
#include <stlsoft/synch/lock_scope.hpp>
#include <winstl/synch/thread_mutex.hpp>
#include <windows.h>
#include <process.h>
#include "process_monitor.hpp"
#include "process.hpp"
#include "exception.hpp"

using namespace std;
using stlsoft::lock_scope;
using winstl::thread_mutex;

namespace judge {

process_monitor::debugger::debugger(process_monitor &pm)
	: debug_object(true), pm_(pm)
{}

NTSTATUS process_monitor::debugger::on_exit_process(PDBGUI_WAIT_STATE_CHANGE state_change)
{
	shared_ptr<context_base> context;
	{
		lock_scope<thread_mutex> lock(pm_.contexts_mutex_);
		auto iter = pm_.contexts_.find(reinterpret_cast<uint32_t>(
			state_change->AppClientId.UniqueProcess));
		if (iter == pm_.contexts_.end())
			return DBG_CONTINUE;
		context = iter->second;
		pm_.contexts_.erase(iter);
	}
	context->set_exit_status(static_cast<int32_t>(
		state_change->StateInfo.ExitProcess.ExitStatus));
	context->invoke();
	return DBG_CONTINUE;
}

NTSTATUS process_monitor::debugger::on_exception(PDBGUI_WAIT_STATE_CHANGE state_change)
{
	EXCEPTION_RECORD &record = state_change->StateInfo.Exception.ExceptionRecord;
	if (!state_change->StateInfo.Exception.FirstChance) {
		shared_ptr<context_base> context;
		{
			lock_scope<thread_mutex> lock(pm_.contexts_mutex_);
			auto iter = pm_.contexts_.find(reinterpret_cast<uint32_t>(
				state_change->AppClientId.UniqueProcess));
			if (iter == pm_.contexts_.end())
				return DBG_EXCEPTION_NOT_HANDLED;
			context = iter->second;
		}
		context->set_exception(shared_ptr<EXCEPTION_RECORD>(new EXCEPTION_RECORD(record)));
		try {
			context->job_object()->terminate(static_cast<int32_t>(record.ExceptionCode));
		} catch (const win32_exception &) {
		}
		return DBG_EXCEPTION_HANDLED;
	} else if (record.ExceptionCode == EXCEPTION_BREAKPOINT) {
		return DBG_EXCEPTION_HANDLED;
	} else {
		return DBG_EXCEPTION_NOT_HANDLED;
	}
}

process_monitor::context_base::context_base(const std::shared_ptr<judge::job_object> &job)
	: job_(job)
{}

void process_monitor::context_base::set_exit_status(int32_t exit_status)
{
	result_.exit_status = exit_status;
}

void process_monitor::context_base::set_exception(const shared_ptr<EXCEPTION_RECORD> &exception)
{
	result_.exception = exception;
}

const shared_ptr<judge::job_object> &process_monitor::context_base::job_object()
{
	return job_;
}

process_monitor::process_monitor()
#pragma warning(disable:4355)
	: debugger_(*this)
	, thread_handle_(_create_thread())
{}

process_monitor::~process_monitor()
{
	debugger_.close();
	::WaitForSingleObject(thread_handle_, INFINITE);
	::CloseHandle(thread_handle_);
}

HANDLE process_monitor::_create_thread()
{
	HANDLE result = reinterpret_cast<HANDLE>(
		::_beginthreadex(nullptr, 0, _thread_entry, this, 0, nullptr));
	if (!result) {
		throw judge_exception(JSTATUS_GENERIC_ERROR);
	}
	return result;
}

unsigned int __stdcall process_monitor::_thread_entry(void *param)
{
	process_monitor &pm = *reinterpret_cast<process_monitor *>(param);
	while (pm.debugger_.wait_for_event())
	{}
	return 0;
}

}
